package com.kh.springwebSocket.config;

import org.springframework.context.annotation.Configuration;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.config.ChannelRegistration;
import org.springframework.messaging.simp.config.MessageBrokerRegistry;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptorAdapter;
import org.springframework.security.authentication.AuthenticationCredentialsNotFoundException;
import org.springframework.security.config.annotation.web.messaging.MessageSecurityMetadataSourceRegistry;
import org.springframework.security.config.annotation.web.socket.AbstractSecurityWebSocketMessageBrokerConfigurer;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.config.annotation.EnableWebSocketMessageBroker;
import org.springframework.web.socket.config.annotation.StompEndpointRegistry;
import org.springframework.web.socket.config.annotation.WebSocketTransportRegistration;
import org.springframework.web.socket.handler.WebSocketHandlerDecorator;

import java.security.Principal;

/**
 * Created by crossice on 2017/5/13.
 */
@Configuration
@EnableWebSocketMessageBroker
public class WebSocketConfig extends AbstractSecurityWebSocketMessageBrokerConfigurer {
    private static final String WEBSOCKET_POINT = "/api-messages";

    private static final String WEBSOCKET_USER = "/user";

    @Override
    public void registerStompEndpoints(StompEndpointRegistry registry) {

        registry.addEndpoint(WEBSOCKET_POINT).setAllowedOrigins("*").withSockJS().setInterceptors();
    }

    @Override
    public void configureWebSocketTransport(WebSocketTransportRegistration registration) {

        //websocket链接记录
        registration.addDecoratorFactory(handler ->
                new WebSocketHandlerDecorator(handler) {
                    @Override
                    public void afterConnectionEstablished(final WebSocketSession session) throws Exception {
                        Principal principal = session.getPrincipal();
                        System.out.println("websocket-online: " + (principal != null ? principal.getName() : null));
                        super.afterConnectionEstablished(session);
                    }

                    @Override
                    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus)
                            throws Exception {
                        Principal principal = session.getPrincipal();
                        System.out.println("websocket-online: " + (principal != null ? principal.getName() : null));
                        super.afterConnectionClosed(session, closeStatus);
                    }

                }
        );
        super.configureWebSocketTransport(registration);

    }

    @Override
    protected void customizeClientInboundChannel(ChannelRegistration registration) {
        registration.setInterceptors(new ChannelInterceptorAdapter() {
            @Override
            public Message<?> preSend(Message<?> message, MessageChannel channel) {
                StompHeaderAccessor accessor = StompHeaderAccessor.wrap(message);

                System.out.println();

                if (StompCommand.SUBSCRIBE.equals(accessor.getCommand())) {

                    String currentLoginUser = SecurityUtils.getCurrentUserLogin();

                    String currentSubscribeUri = accessor.getDestination();
                    //判断当前订阅者是否是
                    if (currentLoginUser != null && currentLoginUser.length() > 0) {
                        Principal principal = accessor.getUser();
                        if (principal == null || (principal != null && currentSubscribeUri != null && !currentSubscribeUri.contains(principal.getName()))) {
                            throw new AuthenticationCredentialsNotFoundException(" current login user is not same with who subscribed");
                        }

                    } else {
                        throw new AuthenticationCredentialsNotFoundException(" need user login");
                    }
                }

                return message;
            }


        });
    }


    @Override
    public void configureMessageBroker(MessageBrokerRegistry registry) {

        registry.setUserDestinationPrefix(WEBSOCKET_USER);
        registry.enableSimpleBroker(WEBSOCKET_USER);
    }

    @Override
    protected boolean sameOriginDisabled() {
        return true;
    }

    @Override
    protected void configureInbound(MessageSecurityMetadataSourceRegistry messages) {
        messages.simpSubscribeDestMatchers("/user/**").authenticated();
    }
}
